import numpy as np
from scipy.fftpack import fft
import matplotlib.pyplot as plt
import cv2 as cv


def myfft_function(signal, begin, end):
    N = end - begin  # 采样点数
    sample_freq = 20  # 采样频率 20 Hz, 大于两倍的最高频率
    sample_interval = 1 / sample_freq  # 采样间隔s
    signal_len = N * sample_interval  # 信号长度
    t = np.arange(0, signal_len, sample_interval)

    fft_data = fft(signal)

    # 计算幅值
    # 在python的计算方式中，fft结果的直接取模和真实信号的幅值不一样。
    # 对于非直流量的频率，直接取模幅值会扩大N/2倍， 所以需要除了N乘以2。
    # 对于直流量的频率(0Hz)，直接取模幅值会扩大N倍，所以需要除了N。
    fft_amp0 = np.array(np.abs(fft_data) / N * 2)  # 用于计算双边谱
    fft_amp0[0] = 0.5 * fft_amp0[0]
    N_2 = int(N / 2)
    fft_amp = fft_amp0[0:N_2]  # 单边谱

    # 计算相位
    phase_spectrum = np.angle(fft_data, deg=True)  # deg=True表示返回的是角度值，否则是弧度值
    # 将相位从-180~180转换到0~255
    n180 = np.full(N, 180)  # 创建含有N个数的数组，值都为180
    phase_spectrum = (phase_spectrum + n180) * 0.7
    fft_pha = phase_spectrum[0:N_2]

    # 计算频谱的频率轴
    list1 = np.array(range(0, int(N / 2)))
    freq = sample_freq * list1 / N  # 单边谱的频率轴

    return freq, fft_amp, fft_pha, t


img_path = r"E:\studylife\detectflaws\code\testcpp\lockPha\lockPha\mapAndResult\images\7_18\40cmhz0.5"  # 原视频裁切出来的图片的存放地址
save_path = r"E:\studylife\detectflaws\code\testcpp\lockPha\lockPha\mapAndResult\fftResult\7_18\40cmhz0.5"

column1 = 300
row1 = 100
dt = 50  # 点间间隔


column2 = column1 + dt
row2 = row1 + dt

column3 = column2 + dt
row3 = row2 + dt

column4 = column3 + dt
row4 = row3 + dt

# 读取的图片帧数
map_begin = 0
map_end = 500


# 创建列表存储对应帧数视频图像
lists_gray = []
array_gray = []

# 将视频存入array_gray中
for k in range(map_begin, map_end):  # 读取map_quantity帧图片计算fft
    img = cv.imread(f'{img_path}/{k}.jpg', flags=0)  # flags=0读取为灰度图像
    # 为列表赋值
    lists_gray.append(img)
array_gray = np.array(lists_gray)  # 把lists_gray转换成np.array类型方便后续裁切

# 存储各信号灰度值的向量
array_gray1 = array_gray[:, row1, column1]
array_gray2 = array_gray[:, row2, column2]
array_gray3 = array_gray[:, row3, column3]
array_gray4 = array_gray[:, row4, column4]

# 对信号进行傅里叶分解求取频率轴，幅度值，相位值，信号时间
freq1, fft_amp1, fft_pha1, t = myfft_function(array_gray1, map_begin, map_end)
freq2, fft_amp2, fft_pha2, t = myfft_function(array_gray2, map_begin, map_end)
freq3, fft_amp3, fft_pha3, t = myfft_function(array_gray3, map_begin, map_end)
freq4, fft_amp4, fft_pha4, t = myfft_function(array_gray4, map_begin, map_end)


# 绘制结果
plt.figure()
# 原信号
plt.subplot(121)
plt.plot(t, array_gray1, 'r')
plt.plot(t, array_gray2, 'y')
plt.plot(t, array_gray3, 'b')
plt.plot(t, array_gray4, 'g')
plt.title(' Original signal')
plt.xlabel('t (s)')
plt.ylabel(' Amplitude ')

# 幅值谱
plt.subplot(122)
plt.plot(freq1, fft_amp1, 'r')
plt.plot(freq2, fft_amp2, 'y')
plt.plot(freq3, fft_amp3, 'b')
plt.plot(freq4, fft_amp4, 'g')
plt.title(' FuPinTu')
# plt.ylim(0, 250)
plt.xlim(0, 1)
plt.xlabel('frequency  (Hz)')
plt.ylabel(' amp ')

# # 相位谱
# plt.subplot(122)
# plt.plot(freq1, fft_pha1, 'r')
# plt.plot(freq2, fft_pha2, 'y')
# plt.plot(freq3, fft_pha3, 'b')
# plt.plot(freq4, fft_pha4, 'g')
# plt.title(' PHA')
# # plt.ylim(0, 250)
# plt.xlim(0, 1)
# plt.xlabel('frequency  (Hz)')
# plt.ylabel(' pha ')

plt.gcf().set_size_inches(18, 6)
plt.savefig(save_path + r"\result.jpg")
plt.show()
